from typing import Optional
from pydantic import BaseModel, Field, ConfigDict
from pydantic.json_schema import JsonSchemaValue
from pydantic import GetJsonSchemaHandler
from bson import ObjectId

# 自定义 ObjectId 类型
class PyObjectId(ObjectId):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        if not ObjectId.is_valid(v):
            raise ValueError("Invalid objectid")
        return ObjectId(v)

    @classmethod
    def __get_pydantic_json_schema__(
        cls, _core_schema: dict, handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        return handler({"type": "string"})  # 定义为字符串类型

# MongoDB 基础模型（处理 _id 和 JSON 序列化）
class MongoBaseModel(BaseModel):
    # 使用 default_factory 自动生成 ObjectId，这样在创建模型实例时如果不提供 _id，会自动生成一个。
    id: Optional[PyObjectId] = Field(default_factory=PyObjectId, alias="_id")  # 映射 MongoDB 的 _id 字段

    model_config = ConfigDict(
        populate_by_name=True,  # 启用字段名填充（替代 allow_population_by_field_name）
        json_encoders={ObjectId: str},  # 序列化 ObjectId 为字符串
        arbitrary_types_allowed=True,  # 允许自定义类型
    )